#include "TFile.h"
#include "TTree.h"
#include "TChain.h"
#include "TF1.h"
#include "TGraph.h"
#include "TMath.h"
#include "TFitResult.h"
#include "TCanvas.h"
#include "TAxis.h"

#include <vector>
#include <string>

#include "LineShapeExtract.hh"

LineShapeExtract::LineShapeExtract(std::string yamlPath)
{
    m_Config = YAML::LoadFile(yamlPath);
}

LineShapeExtract::~LineShapeExtract(){}

std::pair<int, double> LineShapeExtract::FindMinInGraph(TGraph* graph)
{    
    Int_t numOfPoint = graph->GetN();
    Int_t minPoint = -1;
    Double_t minVal = 1e10;
    for(Int_t i = 0; i < numOfPoint; ++i)
    {
        if (minVal > graph->GetPointY(i))
        {
            minPoint = i;
            minVal = graph->GetPointY(i);
        }
    }
    return std::pair<int, double>(minPoint, minVal);
}

Double_t LineShapeCDEX(Double_t* x, Double_t* par)
{
    Double_t expTerm = TMath::Exp(TMath::Power(par[4], 2) / (2 * TMath::Power(par[3], 2)) - (x[0] - par[2]) / par[3]);
    Double_t erfTerm = 1 + TMath::Erf((par[3] * (x[0] - par[2]) - TMath::Power(par[4], 2)) / (TMath::Sqrt(2) * par[3] * par[4]));
    //Double_t val = par[0] - par[1] / (2 * par[3]) * expTerm * erfTerm;
    return (par[0] - par[1] / (2 * par[3]) * expTerm * erfTerm);
}

void LineShapeExtract::Processing()
{
    std::string treeName = m_Config["TreeName"].as<std::string>();
    std::vector<std::string> calibFileList = m_Config["CalibFileList"].as<std::vector<std::string> >();
    std::string outName = m_Config["OutFileName"].as<std::string>();
    Int_t fixTau = m_Config["FixTau"].as<int>();
    Double_t tauVal = m_Config["TauVal"].as<double>();

    Double_t fitRange[2] = {300., 500.};

    TChain* chain = new TChain(treeName.c_str());
    for(auto fileName:calibFileList)
    {
        chain->Add(fileName.c_str());
    }

    TFile* outFile = new TFile(outName.c_str(), "recreate");
    TTree* outTree = chain->CloneTree(0);

    TGraph* sigwave_ch0 = new TGraph();
    Double_t bline_ch0 = 0.;
    Double_t sigv_ch0 = 0.;
    chain->SetBranchAddress("sigwave_ch0",  &sigwave_ch0);
    chain->SetBranchAddress("bline_ch0",    &bline_ch0);
    chain->SetBranchAddress("sigv_ch0",     &sigv_ch0);

    Double_t baseLine = 0.;
    Double_t amp = 0.;
    Double_t time = 0.;
    Double_t tau = 0.;
    Double_t smear = 0.;
    Int_t isValid = 0;
    Double_t fitChi2 = 0.;
    Int_t status = -1;
    outTree->Branch("BaseLine",     &baseLine,      "BaseLine/D");
    outTree->Branch("Amp",          &amp,           "Amp/D");
    outTree->Branch("Time",         &time,          "Time/D");
    outTree->Branch("Tau",          &tau,           "Tau/D");
    outTree->Branch("Smear",        &smear,         "Smear/D");
    outTree->Branch("IsValid",      &isValid,       "IsValid/I");
    outTree->Branch("FitChi2",      &fitChi2,       "FitChi2/D");
    outTree->Branch("Status",       &status,        "Status/I");

    TF1* fitFunc = new TF1("fitFunc", LineShapeCDEX, fitRange[0], fitRange[1], 5);
    fitFunc->SetParLimits(0, 850, 950);
    fitFunc->SetParLimits(1, 0, 8000);
    fitFunc->SetParLimits(2, 350, 450);
    fitFunc->SetParLimits(3, 1., 6.);
    fitFunc->SetParLimits(4, 0., 4.);

    if(fixTau != 0)
    {
        fitFunc->FixParameter(3, tauVal);
    }

    Long64_t numOfEntries = chain->GetEntries();
    for(Long64_t entry = 0; entry < numOfEntries; ++entry)
    {
        chain->GetEntry(entry);
        std::pair<int, double>minInGraph = FindMinInGraph(sigwave_ch0);
        Double_t ampRaw = bline_ch0 - sigv_ch0;
        if(ampRaw < 10 || ampRaw != minInGraph.second)
        {
            continue;
        }

        fitFunc->SetParameters(bline_ch0, sigv_ch0, minInGraph.first, tauVal, 1.);
        TFitResultPtr result = sigwave_ch0->Fit(fitFunc, "QS", "", fitRange[0], fitRange[1]);
        const double* fitVal = result->GetParams();
        baseLine = fitVal[0];
        amp = fitVal[1];
        time = fitVal[2];
        tau = fitVal[3];
        smear = fitVal[4];
        isValid = result->IsValid();
        fitChi2 = result->Chi2();
        status = result->Status();

        outTree->Fill();

        if(entry%1000 == 0)
        {
            std::cout << "Processing: " << entry << " / " << numOfEntries << std::endl;
        }
    }

    outFile->cd();
    outTree->AutoSave();

    delete fitFunc;
    delete chain;
}